package cn.yx.framework.security;

import javax.crypto.*;
import javax.crypto.spec.DESKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;

public class SymmetricalCoder {
    private String mAlgorithm;
    private byte[] mKey;

    public SymmetricalCoder(String algorithm, String secretKey) throws NoSuchAlgorithmException {
        this.mAlgorithm = algorithm;
        this.mKey = this.initKey(algorithm, secretKey.getBytes());
    }

    public byte[] encrypt(byte[] data) throws InvalidKeySpecException, InvalidKeyException, NoSuchAlgorithmException, NoSuchPaddingException, BadPaddingException, IllegalBlockSizeException {
        Key k = this.getKey(this.mAlgorithm, this.mKey);
        Cipher cipher = Cipher.getInstance(this.mAlgorithm);
        cipher.init(1, k);
        return cipher.doFinal(data);
    }

    public byte[] decrypt(byte[] data) throws InvalidKeyException, InvalidKeySpecException, NoSuchAlgorithmException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException {
        Key k = this.getKey(this.mAlgorithm, this.mKey);
        Cipher cipher = Cipher.getInstance(this.mAlgorithm);
        cipher.init(2, k);
        return cipher.doFinal(data);
    }

    public String encrypt(String source) throws NoSuchPaddingException, NoSuchAlgorithmException, IllegalBlockSizeException, BadPaddingException, InvalidKeyException, InvalidKeySpecException {
        return new String(this.encrypt(source.getBytes()));
    }

    private byte[] initKey(String algorithm) throws NoSuchAlgorithmException {
        return this.initKey(algorithm, (byte[])null);
    }

    public byte[] initKey(String algorithm, byte[] seed) throws NoSuchAlgorithmException {
        SecureRandom secureRandom;
        if(seed != null) {
            secureRandom = new SecureRandom(seed);
        } else {
            secureRandom = new SecureRandom();
        }

        KeyGenerator kg = KeyGenerator.getInstance(algorithm);
        kg.init(secureRandom);
        return kg.generateKey().getEncoded();
    }

    private Key getKey(String algorithm, byte[] key) throws InvalidKeyException, InvalidKeySpecException, NoSuchAlgorithmException {
        if("DES".equals(algorithm)) {
            DESKeySpec dks = new DESKeySpec(key);
            SecretKeyFactory keyFactory = SecretKeyFactory.getInstance(algorithm);
            return keyFactory.generateSecret(dks);
        } else {
            return new SecretKeySpec(key, algorithm);
        }
    }

    interface ALGORITHM {
        String DES = "DES";
        String DESEDE = "DESede";
        String AES = "AES";
        String BLOWFISH = "Blowfish";
        String RC2 = "RC2";
        String RC4 = "RC4";
    }
}
